package com.sxd.bill.ner.utils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class SeqEntityScore {
    private List<LabelChunk> origins;
    private List<LabelChunk> founds;
    private List<LabelChunk> rights;

    public SeqEntityScore() {
        this.origins = new ArrayList<>();
        this.founds = new ArrayList<>();
        this.rights = new ArrayList<>();
    }

    public void update(List<String> labels, List<String> preds) {
        List<LabelChunk> labelEntities = OutputUtility.getLabelChunks(labels);
        List<LabelChunk> preEntities = OutputUtility.getLabelChunks(preds);
        List<LabelChunk> commons = labelEntities.stream()
                .filter(item -> preEntities.contains(item))
                .collect(Collectors.toList());
        origins.addAll(labelEntities);
        founds.addAll(preEntities);
        rights.addAll(commons);
    }

    private Map<String, Float> compute(long origin, long found, long right) {
        float recall = origin == 0 ? 0 : (float) right / origin;
        float precision = found == 0 ? 0 : (float) right / found;
        float f1 = recall + precision == 0 ? 0 : (float) (2 * precision * recall) / (precision + recall);
        Map<String, Float> score = new HashMap<>();
        score.put("recall", recall);
        score.put("precision", precision);
        score.put("f1", f1);
        return score;
    }

    public Map<String, Map<String, Float>> getResult() {
        Map<String, Long> originCounter = origins.stream()
                .map(item -> item.getLabelName())
                .collect(Collectors.groupingBy(w -> w, Collectors.counting()));
        Map<String, Long> foundCounter = founds.stream()
                .map(item -> item.getLabelName())
                .collect(Collectors.groupingBy(w -> w, Collectors.counting()));
        Map<String, Long> rightCounter = rights.stream()
                .map(item -> item.getLabelName())
                .collect(Collectors.groupingBy(w -> w, Collectors.counting()));

        Map<String, Map<String, Float>> scores = new HashMap<>();
        for (String key : originCounter.keySet()) {
            long originCnt = originCounter.get(key);
            long foundCnt = foundCounter.getOrDefault(key, (long) 0);
            long rightCnt = rightCounter.getOrDefault(key, (long) 0);
            Map<String, Float> score = compute(originCnt, foundCnt, rightCnt);
            scores.put(key, score);
        }
        Map<String, Float> scoreTotal = compute(origins.size(), founds.size(), rights.size());
        scores.put("total", scoreTotal);
        return scores;
    }

    public void printScores() {
        Map<String, Map<String, Float>> scores = getResult();
        for (String key : scores.keySet()) {
            Map<String, Float> score = scores.get(key);
            float recall = score.get("recall");
            float precision = score.get("precision");
            float f1 = score.get("f1");
            String scoreLine = String.format("%s  ==> precision: %f recall: %f  f1: %f", key, precision, recall, f1);
            System.out.println(scoreLine);
        }
    }

}
